import torch
from torch.utils.data import Dataset, DataLoader
from torchvision import transforms
from PIL import Image
import os
import pandas as pd
import json
class MultiModalDataset(Dataset):
    def __init__(self, txt_file, transform=None):
        """
        txt_file: 包含图片路径和标签的txt文件路径
        transform: 预处理操作（例如：图像缩放、归一化等）
        """
        self.data = []
        self.transform = transform

        # 读取txt文件并解析图片路径和标签
        with open(txt_file, 'r') as f:
            lines = f.readlines()
            for line in lines:
                paths_and_label = line.strip().split(',')
                DIFF_image_path = paths_and_label[0].strip()  # motley 图像路径
                WNB_image_path = paths_and_label[1].strip()  # cyan 图像路径
                label = int(paths_and_label[2].strip())  # 标签（0, 1, 2）

                # 将路径和标签保存为元组
                self.data.append((DIFF_image_path, WNB_image_path, label))

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        DIFF_image_path, WNB_image_path, label = self.data[idx]

        # 读取图片
        DIFF_image = Image.open(DIFF_image_path).convert("RGB")
        WNB_image = Image.open(WNB_image_path).convert("RGB")

        if self.transform:
            DIFF_image = self.transform(DIFF_image)
            WNB_image = self.transform(WNB_image)

        return DIFF_image, WNB_image, label


# 定义图像预处理（例如：缩放、归一化等）
transform = transforms.Compose([
    transforms.Resize((224, 224)),  # 假设你使用ResNet50并且输入尺寸是224x224
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),  # ImageNet的标准化值
])



base_dir=r"D:\competition\2025tongji\new_data"
subfolders = ['000', '001', '010', '100']
xls_folders = ['000相对正常.xlsx','0原始 - 0幼稚粒 - 1异淋.xlsx','0原始 - 1幼稚粒 - 0异淋.xlsx','1原始 - 0幼稚粒 - 0异淋.xlsx']
output_file = 'new_data_train.txt'  # 修改为datas.txt

max_dict = {'\tWBC': 949.27, '\tNeu#': 143.57, '\tLym#': 62.3, '\tMon#': 26.09, '\tEos#': 25.43, '\tBas#': 8.17, '\tIMG#': 56.82, '\tNeu%': 99.6,
            '\tLym%': 99.0, '\tMon%': 76.4, '\tEos%': 65.6, '\tBas%': 26.9, '\tIMG%': 63.5, '\tRBC': 7.6, '\tHGB': 216.0, '\tHCT': 66.2, '\tMCV': 163.3,
            '\tMCH': 80.6, '\tMCHC': 808.0, '\tRDW-CV': 50.1, '\tRDW-SD': 185.1, '\tPLT': 4823.0, '\tMPV': 17.1, '\tPDW': 19.3, '\tPCT': 4.63, '\tP-LCC': 1225.0,
            '\tP-LCR': 76.0, '\tNRBC#': 27.611, '\tNRBC%': 278.42, '\tHFC#': 4.14, '\tHFC%': 44.2, '\tIME%': 4.1, '\tIME#': 0.77, '\tNLR': 96.67, '\tInR#': 0.72,
            '\tInR‰': 0.15, '\tMicro#': 4.17, '\tMicro%': 75.7, '\tMacro#': 2.2, '\tMacro%': 79.6, '\tPLR': 5760.69, '\tPDW-SD': 33.6}

min_dict = {'\tWBC': 0.07, '\tNeu#': 0.0, '\tLym#': 0.0, '\tMon#': 0.0, '\tEos#': 0.0, '\tBas#': 0.0, '\tIMG#': 0.0, '\tNeu%': 0.0, '\tLym%': 0.0, '\tMon%': 0.0,
            '\tEos%': 0.0, '\tBas%': 0.0, '\tIMG%': 0.0, '\tRBC': 0.18, '\tHGB': 4.0, '\tHCT': 1.4, '\tMCV': 45.7, '\tMCH': 14.6, '\tMCHC': 160.0, '\tRDW-CV': 1.9,
            '\tRDW-SD': 2.7, '\tPLT': 0.0, '\tMPV': 6.5, '\tPDW': 14.1, '\tPCT': 0.003, '\tP-LCC': 0.0, '\tP-LCR': 4.8, '\tNRBC#': 0.0, '\tNRBC%': 0.0, '\tHFC#': 0.0,
            '\tHFC%': 0.0, '\tIME%': 0.0, '\tIME#': 0.0, '\tNLR': 0.0, '\tInR#': 0.0, '\tInR‰': 0.0, '\tMicro#': 0.01, '\tMicro%': 0.2, '\tMacro#': 0.0, '\tMacro%': 0.0,
            '\tPLR': 0.0, '\tPDW-SD': 2.7}


# 用于存储每个子文件夹中的图片数量
image_count_per_category = {}

with open(output_file, 'w') as f:
    # 遍历每个子文件夹和标签
    for label,name in enumerate(subfolders):
        image_count = 0
        Wnb_base_path = os.path.join(os.path.join(base_dir,name),"Wnb_images")
        Diff_base_path =os.path.join(os.path.join(base_dir,name),"Diff_images")
        xls_path = os.path.join(base_dir,xls_folders[label])
        xls_db = pd.read_excel(xls_path,'Sheet1')
        print("DataFrame 的列名：", xls_db.columns)
        # 将值为 '****' 的数据替换为 -1
        xls_db = xls_db.replace('****', -1)
        for name in os.listdir(Wnb_base_path):
            # 构建DIFF和WNB图像的路径
            if name.endswith('.png'):
                DIFF_subfolder_path = os.path.join(Diff_base_path, name)
                WNB_subfolder_path = os.path.join(Wnb_base_path, name)
                bianhao= name.split('.')[0][1:]
                a=xls_db['\t样本编号']

                try:

                    # hang = xls_db[xls_db['\t样本编号'] == float(bianhao)]
                    # WBC = (float(hang['\tWBC']) - min_dict['\tWBC']) / (max_dict['\tWBC'] - min_dict['\tWBC'])
                    # Neu = (float(hang['\tNeu#']) - min_dict['\tNeu#']) / (max_dict['\tNeu#'] - min_dict['\tNeu#'])
                    # Lym = (float(hang['\tLym#']) - min_dict['\tLym#']) / (max_dict['\tLym#'] - min_dict['\tLym#'])
                    # Mon = (float(hang['\tMon#']) - min_dict['\tMon#']) / (max_dict['\tMon#'] - min_dict['\tMon#'])
                    # Eos = (float(hang['\tEos#']) - min_dict['\tEos#']) / (max_dict['\tEos#'] - min_dict['\tEos#'])
                    # Bas = (float(hang['\tBas#']) - min_dict['\tBas#']) / (max_dict['\tBas#'] - min_dict['\tBas#'])
                    # IMG = (float(hang['\tIMG#']) - min_dict['\tIMG#']) / (max_dict['\tIMG#'] - min_dict['\tIMG#'])
                    # Neu_ = (float(hang['\tNeu%']) - min_dict['\tNeu%']) / (max_dict['\tNeu%'] - min_dict['\tNeu%'])
                    # Lym_ = (float(hang['\tLym%']) - min_dict['\tLym%']) / (max_dict['\tLym%'] - min_dict['\tLym%'])
                    # Mon_ = (float(hang['\tMon%']) - min_dict['\tMon%']) / (max_dict['\tMon%'] - min_dict['\tMon%'])
                    # Eos_ = (float(hang['\tEos%']) - min_dict['\tEos%']) / (max_dict['\tEos%'] - min_dict['\tEos%'])
                    # Bas_ = (float(hang['\tBas%']) - min_dict['\tBas%']) / (max_dict['\tBas%'] - min_dict['\tBas%'])
                    # IMG_ = (float(hang['\tIMG%']) - min_dict['\tIMG%']) / (max_dict['\tIMG%'] - min_dict['\tIMG%'])
                    # RBC = (float(hang['\tRBC']) - min_dict['\tRBC']) / (max_dict['\tRBC'] - min_dict['\tRBC'])
                    # HGB = (float(hang['\tHGB']) - min_dict['\tHGB']) / (max_dict['\tHGB'] - min_dict['\tHGB'])
                    # HCT = (float(hang['\tHCT']) - min_dict['\tHCT']) / (max_dict['\tHCT'] - min_dict['\tHCT'])
                    # MCV = (float(hang['\tMCV']) - min_dict['\tMCV']) / (max_dict['\tMCV'] - min_dict['\tMCV'])
                    # MCH = (float(hang['\tMCH']) - min_dict['\tMCH']) / (max_dict['\tMCH'] - min_dict['\tMCH'])
                    # MCHC = (float(hang['\tMCHC']) - min_dict['\tMCHC']) / (max_dict['\tMCHC'] - min_dict['\tMCHC'])
                    # CV = (float(hang['\tRDW-CV']) - min_dict['\tRDW-CV']) / (
                    #             max_dict['\tRDW-CV'] - min_dict['\tRDW-CV'])
                    # SD = (float(hang['\tRDW-SD']) - min_dict['\tRDW-SD']) / (
                    #             max_dict['\tRDW-SD'] - min_dict['\tRDW-SD'])
                    # PLT = (float(hang['\tPLT']) - min_dict['\tPLT']) / (max_dict['\tPLT'] - min_dict['\tPLT'])
                    # MPV = (float(hang['\tMPV']) - min_dict['\tMPV']) / (max_dict['\tMPV'] - min_dict['\tMPV'])
                    # PDW = (float(hang['\tPDW']) - min_dict['\tPDW']) / (max_dict['\tPDW'] - min_dict['\tPDW'])
                    # PCT = (float(hang['\tPCT']) - min_dict['\tPCT']) / (max_dict['\tPCT'] - min_dict['\tPCT'])
                    # LCC = (float(hang['\tP-LCC']) - min_dict['\tP-LCC']) / (max_dict['\tP-LCC'] - min_dict['\tP-LCC'])
                    # LCR = (float(hang['\tP-LCR']) - min_dict['\tP-LCR']) / (max_dict['\tP-LCR'] - min_dict['\tP-LCR'])
                    # NRBC = (float(hang['\tNRBC#']) - min_dict['\tNRBC#']) / (max_dict['\tNRBC#'] - min_dict['\tNRBC#'])
                    # NRBC_ = (float(hang['\tNRBC%']) - min_dict['\tNRBC%']) / (max_dict['\tNRBC%'] - min_dict['\tNRBC%'])
                    # HFC = (float(hang['\tHFC#']) - min_dict['\tHFC#']) / (max_dict['\tHFC#'] - min_dict['\tHFC#'])
                    # HFC_ = (float(hang['\tHFC%']) - min_dict['\tHFC%']) / (max_dict['\tHFC%'] - min_dict['\tHFC%'])
                    # IME = (float(hang['\tIME%']) - min_dict['\tIME%']) / (max_dict['\tIME%'] - min_dict['\tIME%'])
                    # IME_ = (float(hang['\tIME#']) - min_dict['\tIME#']) / (max_dict['\tIME#'] - min_dict['\tIME#'])
                    # NLR = (float(hang['\tNLR']) - min_dict['\tNLR']) / (max_dict['\tNLR'] - min_dict['\tNLR'])
                    # InR = (float(hang['\tInR#']) - min_dict['\tInR#']) / (max_dict['\tInR#'] - min_dict['\tInR#'])
                    # InR_ = (float(hang['\tInR‰']) - min_dict['\tInR‰']) / (max_dict['\tInR‰'] - min_dict['\tInR‰'])
                    # Micro = (float(hang['\tMicro#']) - min_dict['\tMicro#']) / (
                    #             max_dict['\tMicro#'] - min_dict['\tMicro#'])
                    # Micro_ = (float(hang['\tMicro%']) - min_dict['\tMicro%']) / (
                    #             max_dict['\tMicro%'] - min_dict['\tMicro%'])
                    # Macro = (float(hang['\tMacro#']) - min_dict['\tMacro#']) / (
                    #             max_dict['\tMacro#'] - min_dict['\tMacro#'])
                    # Macro_ = (float(hang['\tMacro%']) - min_dict['\tMacro%']) / (
                    #             max_dict['\tMacro%'] - min_dict['\tMacro%'])
                    # PLR = (float(hang['\tPLR']) - min_dict['\tPLR']) / (max_dict['\tPLR'] - min_dict['\tPLR'])
                    # PDW_SD = (float(hang['\tPDW-SD']) - min_dict['\tPDW-SD']) / (
                    #             max_dict['\tPDW-SD'] - min_dict['\tPDW-SD'])
                    #
                    # # 写入文件，格式：DIFF_image_path, WNB_image_path, label
                    # f.write(f"{DIFF_subfolder_path},{WNB_subfolder_path},{label},{WBC},{Neu},{Lym},{Mon},{Eos},{Bas},{IMG},{Neu_},{Lym_},{Mon_},{Eos_},{Bas_},{IMG_},"
                    #         f"{RBC},{HGB},{HCT},{MCV},{MCH},{MCHC},{CV},{SD},{PLT},{MPV},{PDW},{PCT},{LCC},{LCR},{NRBC},{NRBC_},{HFC},{HFC_},{IME},{IME_},{NLR},{InR}"
                    #         f"{InR_},{Micro},{Micro_},{Macro},{Macro_},{PLR},{PDW_SD}\n")

                    hang = xls_db[xls_db['\t样本编号'] == float(bianhao)]
                    keys = [
                        '\tWBC', '\tNeu#', '\tLym#', '\tMon#', '\tEos#', '\tBas#', '\tIMG#',
                        '\tNeu%', '\tLym%', '\tMon%', '\tEos%', '\tBas%', '\tIMG%',
                        '\tRBC', '\tHGB', '\tHCT', '\tMCV', '\tMCH', '\tMCHC',
                        '\tRDW-CV', '\tRDW-SD', '\tPLT', '\tMPV', '\tPDW', '\tPCT',
                        '\tP-LCC', '\tP-LCR', '\tNRBC#', '\tNRBC%', '\tHFC#', '\tHFC%',
                        '\tIME%', '\tIME#', '\tNLR', '\tInR#', '\tInR‰', '\tMicro#',
                        '\tMicro%', '\tMacro#', '\tMacro%', '\tPLR', '\tPDW-SD'
                    ]
                    values = [
                        (float(hang[key]) - min_dict[key]) / (max_dict[key] - min_dict[key])
                        for key in keys
                    ]
                    formatted_values = [f"{value:.8f}" for value in values]

                    # 写入文件，格式：DIFF_image_path, WNB_image_path, label
                    f.write(f"{DIFF_subfolder_path},{WNB_subfolder_path},{label},{','.join(formatted_values)}\n")

                    # 增加图片数量
                    image_count += 1
                    print(f"{DIFF_subfolder_path} is true")
                except:
                    print(f"{DIFF_subfolder_path} is false")
        # 记录每个类别中的图片数量
        image_count_per_category[label] = image_count
# 打印每个类别中的图片数量
print("Image count per category:")
for category, count in image_count_per_category.items():
    print(f"{category}: {count} images")